import { NextRequest, NextResponse } from 'next/server'
import { getServerSession } from 'next-auth'
import { authOptions } from '@/lib/auth'
import prisma from '@/lib/prisma'
import { z } from 'zod'

// Zod schema for validation
const gradeInputSchema = z.object({
  stageSubmissionId: z.string(),
  score: z.number().min(0).nullable(),
  feedback: z.string().nullable(),
});

const gradePayloadSchema = z.object({
  grades: z.array(gradeInputSchema),
});

export async function POST(
  req: NextRequest,
  { params }: { params: { id: string; submissionId: string } }
) {
  const session = await getServerSession(authOptions)

  if (!session || (session.user.role !== 'TEACHER' && session.user.role !== 'ADMIN')) {
    return NextResponse.json({ error: 'Forbidden' }, { status: 403 })
  }

  const { id: experimentId, submissionId } = params // submissionId here refers to the ID of the Submission (formal snapshot)

  if (!experimentId || !submissionId) {
    return NextResponse.json({ error: 'Experiment ID and Submission ID are required' }, { status: 400 });
  }

  let payload;
  try {
    const body = await req.json();
    payload = gradePayloadSchema.parse(body);
  } catch (error) {
    if (error instanceof z.ZodError) {
      return NextResponse.json({ error: 'Invalid input data', details: error.errors }, { status: 400 });
    }
    return NextResponse.json({ error: 'Invalid JSON payload' }, { status: 400 });
  }

  try {
    const updatedFormalSubmission = await prisma.$transaction(async (tx) => {
      let totalScore = 0;
      let allStagesGraded = true;

      // First, verify the main Submission record exists and belongs to the experiment
      const formalSubmissionToGrade = await tx.submission.findUnique({
        where: { id: submissionId },
        select: { experimentId: true } // Select only necessary fields for validation
      });

      if (!formalSubmissionToGrade) {
        throw new Error(`Submission with ID ${submissionId} not found.`);
      }
      if (formalSubmissionToGrade.experimentId !== experimentId) {
        throw new Error(`Submission ${submissionId} does not belong to experiment ${experimentId}.`);
      }

      for (const grade of payload.grades) {
        const stageSubmission = await tx.stageSubmission.findUnique({
          where: { id: grade.stageSubmissionId },
          // include: { stage: true } // stage is not directly needed for update, but good for validation if added
        });

        // Crucially, ensure the StageSubmission being graded belongs to the Submission (snapshot) being graded.
        if (!stageSubmission || stageSubmission.submissionId !== submissionId) {
          throw new Error(`Stage submission ${grade.stageSubmissionId} not found or does not belong to the current formal submission snapshot ${submissionId}.`);
        }

        await tx.stageSubmission.update({
          where: { id: grade.stageSubmissionId },
          data: {
            score: grade.score,
            feedback: grade.feedback,
          },
        });

        if (grade.score !== null) {
          totalScore += grade.score;
        } else {
          allStagesGraded = false; // If any stage is not scored, it's not fully graded.
        }
      }
      
      // Update the overall Submission (formal snapshot) status and score
      const finalGradedSubmission = await tx.submission.update({
        where: { id: submissionId }, // id refers to the Submission model's ID
        data: {
          score: totalScore,
          status: allStagesGraded ? 'GRADED' : 'SUBMITTED', // If not all stages graded, keep as SUBMITTED or a new PARTIALLY_GRADED status
          updatedAt: new Date(),
        },
        include: {
            user: { select: { id: true, name: true, email: true } }, // Changed from student to user
            experiment: { select: { id: true, title: true } },
            stageSubmissions: {
                include: { stage: {select: {id:true, title:true, description:true, order:true, type:true, content:true}} },
                orderBy: { stage: { order: 'asc' } }
            }
        }
      });
      return finalGradedSubmission;
    });

    return NextResponse.json({ submission: updatedFormalSubmission, message: 'Grading saved successfully' });

  } catch (error: any) {
    console.error(`Error grading submission ${submissionId} for experiment ${experimentId}:`, error);
    if (error.message.includes('Stage submission') || error.message.includes('Submission with ID') || error.message.includes('does not belong to experiment')) {
        return NextResponse.json({ error: error.message }, { status: 400 });
    }
    return NextResponse.json({ error: 'Internal server error while grading' }, { status: 500 });
  }
} 